package cn.vire.electric.simulation.config.executor;

import com.alibaba.ttl.threadpool.TtlExecutors;
import lombok.extern.slf4j.Slf4j;
import org.slf4j.MDC;
import org.springframework.boot.autoconfigure.task.TaskExecutionAutoConfiguration;
import org.springframework.boot.task.TaskExecutorBuilder;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.task.TaskDecorator;
import org.springframework.scheduling.annotation.AsyncAnnotationBeanPostProcessor;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import java.util.Map;
import java.util.concurrent.Executor;

/**
 */

@Configuration
@Slf4j
public class ExecutorConfig {

    @Bean(name = {TaskExecutionAutoConfiguration.APPLICATION_TASK_EXECUTOR_BEAN_NAME,
            AsyncAnnotationBeanPostProcessor.DEFAULT_TASK_EXECUTOR_BEAN_NAME,
            "executor"})
    public Executor applicationTaskExecutor(TaskExecutorBuilder builder) {
        ThreadPoolTaskExecutor taskExecutor = builder.build();
        taskExecutor.setTaskDecorator(new TaskDecorator() {
            @Override
            public Runnable decorate(Runnable runnable) {
                Map<String, String> copyOfContextMap = MDC.getCopyOfContextMap();
                if (copyOfContextMap != null) {
                    Runnable decorateRunnable = new Runnable() {
                        @Override
                        public void run() {
                            try {
                                MDC.setContextMap(copyOfContextMap);
                                runnable.run();
                            } finally {
                                MDC.clear();
                            }
                        }
                    };
                    return decorateRunnable;
                }
                return runnable;
            }
        });
        taskExecutor.initialize();
        return TtlExecutors.getTtlExecutor(taskExecutor);

    }
}